{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b2fe72-599d-46ed-8285-25e9e0d31ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import pickle\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "\n",
    "from utils import *\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1275092e-02f5-4538-b43d-33295e8d1151",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data set base dir\n",
    "path_base = Path('datasets/Gaze360')\n",
    "path_images = path_base / 'images'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c7d54df-15fe-45cd-bfc1-d2c9c971b985",
   "metadata": {},
   "source": [
    "### train set "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06b70cc3-7c91-473a-942c-ac18e27872b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(str(path_base / 'data_for_model_2/train_gaze_eyes3D_face68.pkl'), 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "data = {d['name']: d for d in data}\n",
    "paths = list(data.keys())\n",
    "print(f\"Number of samples in dataset: {len(data)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07baab32-872f-4b1e-8bde-3448e9e5049c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomly loop through the data\n",
    "k = np.random.randint(len(paths))\n",
    "path = paths[k]\n",
    "\n",
    "res = data[path]\n",
    "image = cv2.imread(str(path_images / path))\n",
    "lms68 = res['face']['xyz68'].astype(np.float32)\n",
    "eyes = {\n",
    "    'left' : (res['eyes']['left']['P'] @ eyel_template_homo.T).T,\n",
    "    'right': (res['eyes']['right']['P'] @ eyer_template_homo.T).T\n",
    "}\n",
    "gaze_vector = res['gaze']['face']['vector']\n",
    "print(f\"Gaze direction: {gaze_vector}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfcb10b6-1736-41e7-8aa1-94aa60e081d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize gaze direction\n",
    "image_gaze = draw_gaze_from_vector(image.copy(), lms68, gaze_vector, colour=[255, 0, 0])\n",
    "plt.imshow(image_gaze[:, :, [2, 1, 0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb35015f-3e74-4f6b-a812-4424e03fd6eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize 3D eyes\n",
    "image_eyes = draw_eyes(image.copy(), lms68, eyes, colour=[178, 255, 102])\n",
    "plt.imshow(image_eyes[:, :, [2, 1, 0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78fd3619-ef5e-4454-a631-64e1e346ea4c",
   "metadata": {},
   "source": [
    "### test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb9ecad8-c634-4673-ade7-74abfd64d037",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(str(path_base / 'data_for_model_2/test_gaze_eyes3D_face68.pkl'), 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "data = {d['name']: d for d in data}\n",
    "paths = list(data.keys())\n",
    "print(f\"Number of samples in dataset: {len(data)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e98b5f22-2e7e-4b56-8a4f-e83af3aed82c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomly loop through the data\n",
    "k = np.random.randint(len(paths))\n",
    "path = paths[k]\n",
    "\n",
    "res = data[path]\n",
    "image = cv2.imread(str(path_images / path))\n",
    "lms68 = res['face']['xyz68'].astype(np.float32)\n",
    "eyes = {\n",
    "    'left' : (res['eyes']['left']['P'] @ eyel_template_homo.T).T,\n",
    "    'right': (res['eyes']['right']['P'] @ eyer_template_homo.T).T\n",
    "}\n",
    "gaze_vector = res['gaze']['face']['vector']\n",
    "print(f\"Gaze direction: {gaze_vector}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "846fe12b-f370-49d5-a675-b9c2e8482871",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize gaze direction\n",
    "image_gaze = draw_gaze_from_vector(image.copy(), lms68, gaze_vector, colour=[255, 0, 0])\n",
    "plt.imshow(image_gaze[:, :, [2, 1, 0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fa38b79-f840-4a11-b722-a7ff0e4f7555",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize 3D eyes\n",
    "image_eyes = draw_eyes(image.copy(), lms68, eyes, colour=[178, 255, 102])\n",
    "plt.imshow(image_eyes[:, :, [2, 1, 0]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
